{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pprint import pprint\n",
    "import copy\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../../minecraft/python/craftassist/text_to_tree_tool/turk_data/final_prompts_and_humanbot/all_prompts_final.json') as f:\n",
    "    data =  json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4532\n",
      "[['build a wall around the house'], {'dialogue_type': 'HUMAN_GIVE_COMMAND', 'action_sequence': [{'action_type': 'BUILD', 'location': {'location_type': 'REFERENCE_OBJECT', 'relative_direction': 'AROUND', 'reference_object': {'has_name': [0, [5, 5]]}}, 'schematic': {'has_name': [0, [2, 2]]}}]}]\n"
     ]
    }
   ],
   "source": [
    "print(len(data))\n",
    "print(data[50])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../old_data/all_data_newdct.json') as f:\n",
    "    data = json.load(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "from spacy.lang.en import English\n",
    "tokenizer = English().Defaults.create_tokenizer()\n",
    "\n",
    "def tokenize(st):\n",
    "    return ' '.join([str(x) for x in tokenizer(st)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['kindly proceed forty six paces to the left of the yard'], {'dialogue_type': 'HUMAN_GIVE_COMMAND', 'action_sequence': [{'location': {'reference_object': {'has_name': [0, [10, 10]]}, 'steps': [0, [2, 3]], 'relative_direction': 'LEFT', 'location_type': 'REFERENCE_OBJECT'}, 'action_type': 'MOVE'}]}]\n"
     ]
    }
   ],
   "source": [
    "print(data.keys())\n",
    "# do also for templated\n",
    "train_data, valid_data, test_data = data['train'], data['valid'], data['test']\n",
    "print(train_data.keys())\n",
    "print(valid_data.keys())\n",
    "print(test_data.keys())\n",
    "total = {}\n",
    "for data_type in [train_data, valid_data, test_data]:\n",
    "    for key in ['templated', 'rephrased', 'prompts', 'humanbot']:\n",
    "        if key in data_type:\n",
    "            total[key] = total.get(key, 0) + len(data_type[key])\n",
    "print(total)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_snake_case(word, case=\"lower\"):\n",
    "    \"\"\"convert a given word to snake case\"\"\"\n",
    "    s1 = re.sub(\"(.)([A-Z][a-z]+)\", r\"\\1_\\2\", word)\n",
    "    snake_case = re.sub(\"([a-z0-9])([A-Z])\", r\"\\1_\\2\", s1)\n",
    "    if case == \"upper\":\n",
    "        return snake_case.upper()\n",
    "    return snake_case.lower()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_keys(d):\n",
    "    for key, val in d.items():\n",
    "        if key in ['action_reference_object', 'location_reference_object']:\n",
    "            d['reference_object'] = d[key]\n",
    "            del d[key]\n",
    "        elif key.endswith('_'):\n",
    "            d[key[:-1]] = d[key]\n",
    "            del d[key]\n",
    "        if type(val) == dict:\n",
    "            fix_keys(val)\n",
    "        \n",
    "    return d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_vals(d):\n",
    "    for key, val in d.items():\n",
    "        if type(val) == dict:\n",
    "            fix_vals(val)\n",
    "        # fix spans\n",
    "        elif type(val) == list:\n",
    "            d[key] = [0, val]\n",
    "        else:\n",
    "            d[key] = to_snake_case(val, \"upper\")\n",
    "    return d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_loc_type(d):\n",
    "    for key, val in d.items():\n",
    "        if key == \"location_type\" and val in ['SpeakerLook', 'AgentPos', 'SpeakerPos']:\n",
    "            return True\n",
    "        if type(val) == dict:\n",
    "            if check_loc_type(val):\n",
    "                return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "COLOURS = [\"red\", \"blue\", \"green\", \"yellow\", \"purple\", \"grey\", \"brown\", \"black\", \"orange\"]\n",
    "ABSTRACT_SIZE = [\n",
    "    \"really tiny\",\n",
    "    \"very tiny\",\n",
    "    \"tiny\",\n",
    "    \"really small\",\n",
    "    \"very small\",\n",
    "    \"really little\",\n",
    "    \"very little\",\n",
    "    \"small\",\n",
    "    \"little\",\n",
    "    \"medium\",\n",
    "    \"medium sized\",\n",
    "    \"big\",\n",
    "    \"large\",\n",
    "    \"really big\",\n",
    "    \"very big\",\n",
    "    \"really large\",\n",
    "    \"very large\",\n",
    "    \"huge\",\n",
    "    \"gigantic\",\n",
    "    \"really huge\",\n",
    "    \"very huge\",\n",
    "    \"really gigantic\",\n",
    "    \"very gigantic\",\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_new_data(datalist):\n",
    "    new_data = []\n",
    "\n",
    "    skip_data = []\n",
    "\n",
    "    for data_point in datalist:\n",
    "        text = data_point[0]\n",
    "        text = tokenize(text)\n",
    "        re_dict = data_point[1]\n",
    "        \n",
    "        # skip the ones that need re-annotation with coref-type\n",
    "        if check_loc_type(re_dict):\n",
    "            skip_data += [([text], re_dict)]\n",
    "        else:\n",
    "            x = copy.deepcopy(re_dict)\n",
    "            new_dict = fix_keys(x)\n",
    "            n_dict = fix_vals(new_dict)\n",
    "            updated_dict = {}\n",
    "\n",
    "            if n_dict['action_type'] == 'TAG':\n",
    "                updated_dict['dialogue_type'] = 'PUT_MEMORY'\n",
    "                if 'reference_object' in n_dict:\n",
    "                    updated_dict['filters'] = {'reference_object' : n_dict['reference_object']}\n",
    "                if 'tag' in n_dict:\n",
    "                    updated_dict['upsert'] = {'memory_data' : {\"memory_type\": 'TRIPLE'}}\n",
    "                    tag_range = n_dict['tag'][1]\n",
    "                    words = \" \".join(text.split()[tag_range[0]: tag_range[1]+1])\n",
    "                    if words.strip() in COLOURS:\n",
    "                        updated_dict['upsert']['memory_data']['has_colour'] = n_dict['tag']\n",
    "                    elif words.strip() in ABSTRACT_SIZE:\n",
    "                        updated_dict['upsert']['memory_data']['has_size'] = n_dict['tag']\n",
    "                    else:\n",
    "                        updated_dict['upsert']['memory_data']['has_tag'] = n_dict['tag']\n",
    "            elif n_dict['action_type'] == 'ANSWER':\n",
    "                updated_dict['dialogue_type'] = 'GET_MEMORY'\n",
    "                if 'reference_object' not in n_dict or 'answer_type' not in n_dict:\n",
    "                    skip_data.append([[text], re_dict])\n",
    "                    continue\n",
    "                updated_dict['filters'] = {'reference_object' : n_dict['reference_object'], 'type' : 'REFERENCE_OBJECT'}\n",
    "                if n_dict['answer_type'] == 'QUERY_TAG':\n",
    "                    if 'query_tag_name' in n_dict and 'query_tag_val' not in n_dict:\n",
    "                        updated_dict['answer_type'] = 'TAG'\n",
    "                        updated_dict['tag_name'] = 'has_' + n_dict['query_tag_name'].lower()\n",
    "                    elif 'query_tag_name' in n_dict and 'query_tag_val' in n_dict:\n",
    "                        updated_dict['answer_type'] = 'EXISTS'\n",
    "                        val = n_dict['query_tag_val']\n",
    "                        key = 'has_' + n_dict['query_tag_name'].lower()\n",
    "                        updated_dict['filters']['reference_object'][key] = val\n",
    "                    else:\n",
    "                        skip_data.append([[text], re_dict])\n",
    "                        continue\n",
    "                elif n_dict['answer_type'] == 'QUERY_ALL_TAGS':\n",
    "                    updated_dict['answer_type'] = 'TAG'\n",
    "                    updated_dict['tag_name'] = 'has_name'\n",
    "            elif n_dict['action_type'] == 'NOOP':\n",
    "                updated_dict['dialogue_type'] = 'NOOP'\n",
    "            else:\n",
    "                updated_dict[\"dialogue_type\"] = \"HUMAN_GIVE_COMMAND\"\n",
    "                updated_dict[\"action\"] = n_dict\n",
    "            \n",
    "            new_data+= [([text], updated_dict)]\n",
    "\n",
    "    return new_data, skip_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35200\n",
      "25402\n",
      "9798\n"
     ]
    }
   ],
   "source": [
    "# for rephrases\n",
    "rephrases = []\n",
    "new_rephrases, skip_rephrases = [], []\n",
    "rephrases = train_data['rephrased']\n",
    "rephrases.extend(valid_data['rephrased'])\n",
    "rephrases.extend(test_data['rephrased'])\n",
    "\n",
    "# for x in rephrases:\n",
    "#     if 'Name each very large item in front of the orange airport with big' in x[0]:\n",
    "#         print(x)\n",
    "new_rephrases, skip_rephrases = make_new_data(rephrases)\n",
    "print(len(rephrases))\n",
    "print(len(new_rephrases))\n",
    "print(len(skip_rephrases))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(['what are the dimensions of the structure'], {'reference_object': {'location': {'location_type': 'SpeakerLook'}}, 'answer_type': 'query_tag', 'query_tag_name': 'size', 'action_type': 'Answer'})\n"
     ]
    }
   ],
   "source": [
    "print(skip_rephrases[700])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write things that need to be annotated\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/release_data_dialogue/rephrases.json', 'w') as f:\n",
    "    json.dump(new_rephrases, f)\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/reannotate_data/rephrases.json', 'w') as f:\n",
    "    json.dump(skip_rephrases, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3065\n",
      "2513\n",
      "552\n"
     ]
    }
   ],
   "source": [
    "prompts = []\n",
    "prompts = valid_data['prompts']\n",
    "prompts.extend(test_data['prompts'])\n",
    "\n",
    "new_prompts, skip_prompts = make_new_data(prompts)\n",
    "print(len(prompts))\n",
    "print(len(new_prompts))\n",
    "print(len(skip_prompts))\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/release_data_dialogue/prompts.json', 'w') as f:\n",
    "    json.dump(new_prompts, f)\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/reannotate_data/prompts.json', 'w') as f:\n",
    "    json.dump(skip_prompts, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "817 708 109\n"
     ]
    }
   ],
   "source": [
    "# humanbot\n",
    "humanbot = []\n",
    "humanbot = valid_data['humanbot']\n",
    "humanbot.extend(test_data['humanbot'])\n",
    "new_humanbot, skip_humanbot = make_new_data(humanbot)\n",
    "print(len(humanbot), len(new_humanbot), len(skip_humanbot))\n",
    "with open('/checkpoint/kavyasrinet/release_data_dialogue/humanbot.json', 'w') as f:\n",
    "    json.dump(new_humanbot, f)\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/reannotate_data/humanbot.json', 'w') as f:\n",
    "    json.dump(skip_humanbot, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('dialogue_data.json') as f:\n",
    "    dialogue_data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "800000\n",
      "[['please reproduce 90 red porches over here'], {'dialogue_type': 'HUMAN_GIVE_COMMAND', 'action': {'reference_object': {'repeat': {'repeat_key': 'FOR', 'repeat_count': [0, [2, 2]]}, 'has_colour': [0, [3, 3]], 'has_name': [0, [4, 4]], 'has_attribute': [[0, [3, 3]], [0, [4, 4]]]}, 'location': {'coref_resolve': [0, [6, 6]]}, 'action_type': 'BUILD'}}]\n"
     ]
    }
   ],
   "source": [
    "templated = []\n",
    "for key in ['train', 'valid', 'test']:\n",
    "    data = dialogue_data[key]['templated']\n",
    "    templated.extend(data)\n",
    "print(len(templated))\n",
    "print(templated[17])\n",
    "\n",
    "with open('/checkpoint/kavyasrinet/release_data_dialogue/generated_dialogues.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5000\n",
      "790000\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'NOOP'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n",
      "[['what are you doing'], {'dialogue_type': 'GET_MEMORY', 'filters': {'temporal': 'CURRENT', 'type': 'ACTION'}, 'answer_type': 'TAG', 'tag_name': 'action_name'}]\n"
     ]
    }
   ],
   "source": [
    "# print(dialogue_data['test']['templated'][20])\n",
    "print(len(dialogue_data['test']['templated']))\n",
    "print(len(dialogue_data['train']['templated']))\n",
    "for d_type in ['train', 'test', 'valid']:\n",
    "    for example in dialogue_data[d_type]['templated']:\n",
    "        if 'what are you doing' in example[0]:\n",
    "            print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['templated'])\n"
     ]
    }
   ],
   "source": [
    "print(dialogue_data['train'].keys())\n",
    "\n",
    "for x, y in zip(['rephrased'], [new_rephrases]): # add rephrases only\n",
    "# for x, y in zip(['rephrased', 'prompts', 'humanbot'], [new_rephrases, new_prompts, new_humanbot]): \n",
    "    n_train = (9 * len(y)) // 10\n",
    "    n_valid = len(y) // 20\n",
    "    dialogue_data['train'][x] = y[ : n_train]\n",
    "    dialogue_data['valid'][x] = y[n_train : n_train + n_valid]\n",
    "    dialogue_data['test'][x] = y[n_train + n_valid : ]\n",
    "    #print(x, len(dialogue_data['train'][x]), len(dialogue_data['valid'][x]), len(dialogue_data['test'][x]))\n",
    "# write as train, valid, test and each is a dict -> rephrased, prompt, humanbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(['Complete the construction of a couple tiny stairways please'], {'dialogue_type': 'HUMAN_GIVE_COMMAND', 'action': {'reference_object': {'repeat': {'repeat_key': 'FOR', 'repeat_count': [0, [4, 5]]}, 'has_size': [0, [6, 6]], 'has_name': [0, [7, 7]]}, 'action_type': 'FREE_BUILD'}})\n"
     ]
    }
   ],
   "source": [
    "print(dialogue_data['test']['rephrased'][29])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['train', 'valid', 'test'])\n",
      "dict_keys(['templated', 'rephrased'])\n",
      "dict_keys(['templated', 'rephrased'])\n",
      "dict_keys(['templated', 'rephrased'])\n",
      "25402\n",
      "2513\n",
      "708\n",
      "templated 790000\n",
      "rephrased 22861\n",
      "-----------------------------\n",
      "templated 5000\n",
      "rephrased 1270\n",
      "-----------------------------\n",
      "templated 5000\n",
      "rephrased 1271\n",
      "-----------------------------\n"
     ]
    }
   ],
   "source": [
    "print(dialogue_data.keys())\n",
    "# do also for templated\n",
    "train_data, valid_data, test_data = dialogue_data['train'], dialogue_data['valid'], dialogue_data['test']\n",
    "print(train_data.keys())\n",
    "print(valid_data.keys())\n",
    "print(test_data.keys())\n",
    "total = {}\n",
    "print(len(new_rephrases))\n",
    "print(len(new_prompts))\n",
    "print(len(new_humanbot))\n",
    "for data_type in [train_data, valid_data, test_data]:\n",
    "    for key in ['templated', 'rephrased']:\n",
    "        \n",
    "#     for key in ['templated', 'rephrased', 'prompts', 'humanbot']:\n",
    "        print(key, len(data_type[key]))\n",
    "    print(\"-----------------------------\")\n",
    "#         if key in data_type:\n",
    "#             print(key, data_type[key])\n",
    "#             total[key] = total.get(key, 0) + len(data_type[key])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('combined_dialogue_data.json', 'w') as f:\n",
    "    json.dump(dialogue_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('combined_dialogue_data.json') as f:\n",
    "    data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "for data_type in ['train', 'valid', 'test']:\n",
    "    for example in data[data_type]['rephrased']:\n",
    "        if 'good job' in example[0]:\n",
    "            print(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/Users/kavyasrinet/rephrases.json') as f:\n",
    "    data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'list'>\n"
     ]
    }
   ],
   "source": [
    "print(type(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_coref(d):\n",
    "    for key, val in d.items():\n",
    "        if key == \"coref_resolve\":\n",
    "            return True\n",
    "        if type(val) == dict:\n",
    "            if check_coref(val):\n",
    "                return True\n",
    "        if type(val) == list and type(val[0]) == dict:\n",
    "            for v in val:\n",
    "                if check_coref(v):\n",
    "                    return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def location_type(d):\n",
    "    \n",
    "    for key, val in d.items():\n",
    "        #print(key, val)\n",
    "        if key == \"location_type\" and val in ['SPEAKER_LOOK', 'AGENT_POS', 'SPEAKER_POS']:\n",
    "            return True\n",
    "        if type(val) == dict:\n",
    "            if location_type(val):\n",
    "                return True\n",
    "        if type(val) == list and type(val[0]) == dict:\n",
    "            for v in val:\n",
    "                if location_type(v):\n",
    "                    return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35192\n"
     ]
    }
   ],
   "source": [
    "print(len(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_coref(d):\n",
    "    updated_d = {}\n",
    "    for key, val in d.items():\n",
    "        if key == \"coref_resolve\":\n",
    "            updated_d['contains_coreference'] = \"yes\"\n",
    "        elif type(val) == dict:\n",
    "            updated_d[key] = update_coref(val)\n",
    "        elif type(val) == list and type(val[0]) == dict:\n",
    "            updated_d[key] = []\n",
    "            for v in val:\n",
    "                updated_d[key].append(update_coref(v))\n",
    "        else:\n",
    "            updated_d[key] = val\n",
    "    return updated_d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_freebuild(d):\n",
    "    updated_d = {}\n",
    "    for key, val in d.items():\n",
    "        if key == \"action_type\" and val == 'FREE_BUILD':\n",
    "            updated_d[key] = \"FREEBUILD\"\n",
    "        elif type(val) == dict:\n",
    "            updated_d[key] = fix_freebuild(val)\n",
    "        elif type(val) == list and type(val[0]) == dict:\n",
    "            updated_d[key] = []\n",
    "            for v in val:\n",
    "                updated_d[key].append(fix_freebuild(v))\n",
    "        else:\n",
    "            updated_d[key] = val\n",
    "    return updated_d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_action_type(d):\n",
    "    if 'action_sequence' in ad:\n",
    "        return ad['action_sequence'][0]['action_type']\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'DESTROY', 'STOP', 'UNDO', 'MOVE', 'DIG', 'FREE_BUILD', 'BUILD', 'FILL', 'RESUME'}\n"
     ]
    }
   ],
   "source": [
    "action_types = set()\n",
    "for sentence, ad in data:\n",
    "    if 'action_sequence' in ad:\n",
    "        val = ad['action_sequence']\n",
    "        if len(val)> 1:\n",
    "            print(\"YESSSSS\")\n",
    "        for v in val:\n",
    "#             if v['action_type'] == 'FREE_BUILD':\n",
    "#                 print(sentence, ad)\n",
    "            action_types.add(v['action_type'])\n",
    "print(action_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_dig(ad):\n",
    "    new_dict = {}\n",
    "    action_dict = ad['action_sequence'][0]\n",
    "    for key, val in action_dict.items():\n",
    "        if key.startswith('has_'):\n",
    "            if 'schematic' not in new_dict:\n",
    "                new_dict['schematic'] = {}\n",
    "            new_dict['schematic'][key] = val\n",
    "        else:\n",
    "            new_dict[key] = val\n",
    "    out= {}\n",
    "    out['dialogue_type'] = 'HUMAN_GIVE_COMMAND'\n",
    "    out['action_sequence'] = [new_dict]\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_target_action_type(d):\n",
    "    for k, v in d.items():\n",
    "        if k =='contains_coreference':\n",
    "            return True\n",
    "        if type(v) == dict:\n",
    "            if check_target_action_type(v):\n",
    "                return True\n",
    "        if type(v) == list and type(v[0]) == dict:\n",
    "            for val in v:\n",
    "                if check_target_action_type(val):\n",
    "                    return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'dialogue_type': 'HUMAN_GIVE_COMMAND', 'action_sequence': [{'location': {'reference_object': {'has_name': [0, [10, 10]]}, 'steps': [0, [2, 3]], 'relative_direction': 'LEFT', 'location_type': 'REFERENCE_OBJECT'}, 'action_type': 'MOVE'}]}\n",
      "0 0 0 0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "j = 0\n",
    "target = 0\n",
    "updated_data = []\n",
    "dig, fill, freebuild = 0, 0, 0\n",
    "for sentence, ad in data:\n",
    "    \n",
    "    text = sentence\n",
    "    new_d = ad\n",
    "    \n",
    "#     if check_target_action_type(ad):\n",
    "#         print(sentence)\n",
    "#         target += 1\n",
    "    if get_action_type(ad) == 'DIG':\n",
    "        dig += 1\n",
    "        new_d = fix_dig(ad)\n",
    "#     if get_action_type(ad) == 'FILL':\n",
    "#         fill += 1\n",
    "#         print(sentence, ad)\n",
    "#         continue\n",
    "\n",
    "    if get_action_type(new_d) == 'FREE_BUILD':\n",
    "        freebuild += 1\n",
    "        new_d = fix_freebuild(new_d)\n",
    "        \n",
    "#     if location_type(ad):\n",
    "#         i +=1\n",
    "    \n",
    "    # contains coreference\n",
    "    if check_coref(ad):\n",
    "        new_d = update_coref(new_d)\n",
    "    print(new_d)\n",
    "    break\n",
    "    updated_data+= [(text, new_d)]\n",
    "\n",
    "print(i, dig, fill, freebuild)\n",
    "print(target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31870\n"
     ]
    }
   ],
   "source": [
    "# get rephrases data\n",
    "folder = '/Users/kavyasrinet/Desktop/annotated_data/'\n",
    "with open(folder + 'train/rephrases.json') as f:\n",
    "    rephrases = json.load(f)\n",
    "with open(folder + 'valid/rephrases.json') as f:\n",
    "    x = json.load(f)\n",
    "    rephrases.extend(x)\n",
    "with open(folder + 'test/rephrases.json') as f:\n",
    "    x = json.load(f)\n",
    "    rephrases.extend(x)\n",
    "print(len(rephrases))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def location_type(d):\n",
    "    \n",
    "    for key, val in d.items():\n",
    "        #print(key, val)\n",
    "        if key == \"location_type\" and val in ['SPEAKER_LOOK', 'AGENT_POS', 'SPEAKER_POS']:\n",
    "            return True\n",
    "        if type(val) == dict:\n",
    "            if location_type(val):\n",
    "                return True\n",
    "        if type(val) == list and type(val[0]) == dict:\n",
    "            for v in val:\n",
    "                if location_type(v):\n",
    "                    return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n"
     ]
    }
   ],
   "source": [
    "speaker_look = []\n",
    "others = []\n",
    "cnt  = 0\n",
    "for r in rephrases:\n",
    "    sent, ad = r\n",
    "    if location_type(ad):\n",
    "        speaker_look.append(r)\n",
    "    else:\n",
    "        others.append(r)\n",
    "print(cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2359\n",
      "29511\n"
     ]
    }
   ],
   "source": [
    "print(len(speaker_look))\n",
    "print(len(others))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/Users/kavyasrinet/Desktop/speaker_look.txt', 'w') as f:\n",
    "    for x in speaker_look:\n",
    "        sent, tree = x\n",
    "        f.write(sent+\"\\n\")\n",
    "        f.write(json.dumps(tree) + \"\\n\")\n",
    "        f.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2359\n"
     ]
    }
   ],
   "source": [
    "new_r = []\n",
    "with open('/Users/kavyasrinet/Desktop/speaker_look.txt') as f:\n",
    "    i = 0\n",
    "    text = None\n",
    "    ad = None\n",
    "    for line in f.readlines():\n",
    "        line = line.strip()\n",
    "        i += 1\n",
    "        if len(line) == 0:\n",
    "            new_r.append([text, ad])\n",
    "            text = None\n",
    "            ad = None\n",
    "        elif line.startswith('{'):\n",
    "            ad = json.loads(line)\n",
    "        else:\n",
    "            text = line\n",
    "            \n",
    "print(len(new_r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31870\n"
     ]
    }
   ],
   "source": [
    "others.extend(new_r)\n",
    "print(len(others))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_keys(d):\n",
    "    for key, val in d.items():\n",
    "        if key.endswith('_'):\n",
    "            return True\n",
    "        if type(val) == dict:\n",
    "            if find_keys(val):\n",
    "                return True\n",
    "        elif type(val) == list and type(val[0]) ==dict:\n",
    "            for v in val:\n",
    "                if find_keys(v):\n",
    "                    return True\n",
    "        \n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_keys(d):\n",
    "    ad = {}\n",
    "    for key, val in d.items():\n",
    "        if key.endswith('_'):\n",
    "            ad[key[:-1]] = val\n",
    "        elif type(val) == dict:\n",
    "            ad[key]  = fix_keys(val)\n",
    "        elif type(val) == list and type(val[0]) ==dict:\n",
    "            ad[key] = []\n",
    "            for v in val:\n",
    "                ad[key].append(fix_keys(v))\n",
    "        else:\n",
    "            ad[key] = val\n",
    "    return ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "updated_others = []\n",
    "i =0\n",
    "for item in others:\n",
    "    sent, ad = item\n",
    "    if find_keys(ad):\n",
    "        new_ad = fix_keys(ad)\n",
    "        updated_others.append([sent, new_ad])\n",
    "    else:\n",
    "        updated_others.append([sent, ad])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31870\n"
     ]
    }
   ],
   "source": [
    "print(len(updated_others))\n",
    "for item in updated_others:\n",
    "    sent, ad = item\n",
    "    if find_keys(ad):\n",
    "        print(ad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/Users/kavyasrinet/Github/minecraft/python/craftassist/ttad/rephrases_data/rephrases.json', 'w') as f:\n",
    "    json.dump(updated_others, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/Users/kavyasrinet/Github/minecraft/python/craftassist/ttad/rephrases_data/rephrases.json') as f:\n",
    "    data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dict'>\n"
     ]
    }
   ],
   "source": [
    "print(type(data[30][1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
